package com.haibo.config;

import com.zaxxer.hikari.HikariDataSource;
import org.springframework.boot.context.properties.EnableConfigurationProperties;
import org.springframework.context.annotation.Bean;
import org.springframework.context.annotation.Configuration;
import org.springframework.context.annotation.Primary;

import javax.sql.DataSource;
import java.util.Collections;
import java.util.Map;
import java.util.Optional;
import java.util.concurrent.atomic.AtomicReference;
import java.util.stream.Collectors;

@Configuration
@EnableConfigurationProperties(DataSourceProperties.class)
public class DataSourceConfig {

    @Bean
    @Primary
    public DynamicDataSource dynamicDataSource(DataSourceProperties properties) {
        AtomicReference<DataSource> firstDatasource = new AtomicReference<>();
        AtomicReference<DataSource> primaryDatasource = new AtomicReference<>();
        Map<Object, Object> dataSourceMap = Optional.ofNullable(properties.getDynamicSouces())
                .map(ds -> ds.entrySet().stream().collect(Collectors.toMap(e -> (Object) e.getKey(), e -> {
                    DataSourceProperties.DataSourceProp prop = e.getValue();
                    prop.setInitializationFailTimeout(-1);
                    HikariDataSource dataSource = new HikariDataSource(prop);

                    if (firstDatasource.get() == null) {
                        firstDatasource.set(dataSource);
                    }
                    if (prop.isPrimary()) {
                        primaryDatasource.set(dataSource);
                    }
                    return (Object) dataSource;
                }))).orElse(Collections.emptyMap());

        DataSource defaultDataSource = primaryDatasource.get() != null ? primaryDatasource.get() : firstDatasource.get();
        DynamicDataSource dynamicDataSource = new DynamicDataSource();
        dynamicDataSource.setTargetDataSources(dataSourceMap);
        dynamicDataSource.setDefaultTargetDataSource(defaultDataSource);
        return dynamicDataSource;
    }

}
